package myutil.mymap;

import java.util.Comparator;
import java.util.Iterator;
import myutil.mylist.MyLinkedList;
import myutil.mylist.MyList;
import myutil.mylist.MyNoSuchElementException;
import myutil.mylist.MyStack;
import myutil.mymap.MyMap.Entry;
import static myutil.mymap.MyTreeMap.SimpleEntry.BLACK;
import static myutil.mymap.MyTreeMap.SimpleEntry.RED;

/**
 * Implementation of MyMap based on Red-Black tree.
 *
 * @author Hlib_Babii
 * @param <K> type of the key
 * @param <V> type of the value
 */
public class MyTreeMap<K, V> implements MyMap<K, V> {

    private SimpleEntry<K, V> root = null;

    private int size = 0;

    private Comparator comparator = null;

    /**
     * constructs a new, empty tree map, using the natural ordering of its keys.
     */
    public MyTreeMap() {

    }

    /**
     * constructs a new, empty tree map, ordered according to the given
     * comparator.
     *
     * @param comparator comparator according to which elements are ordered
     */
    public MyTreeMap(Comparator comparator) {
        this.comparator = comparator;
    }

    /**
     * Associates the specified value with the specified key in this map.
     *
     * @param key key to be associated with {@code value}
     * @param value value to be associated with {@code key}
     * @return Old value that was associated with key or null if there were no
     * value associated with the specified key
     */
    @Override
    public V put(K key, V value) {
        return insert(key, value, root);
    }

    /**
     * Removes mapping that is associated with the specified key from this map
     *
     * @param key key for the mapping to be removed
     * @return value that was associated with the specified key
     */
    @Override
    public V remove(K key) {
        final SimpleEntry<K, V> entry = getEntry(key);
        return (entry != null) ? removeEntry(entry) : null;
    }

    /**
     * Compares two specified objects with {@code comparator} or
     * {@code Comparable#compareTo} if {@code comparator} is {@code null}.
     *
     * @param obj1 the first objects to be compared
     * @param obj2 the second object to be compared
     * @return a negative integer, zero, or a positive integer as the first
     * argument is less than, equal to, or greater than the second
     * @throws ClassCastException if either comparator is not specified, i.e.
     * {@code comparator == null}, or {@code obj1} or {@code obj2} doesn't
     * implement {@code Comparable}
     */
    private int compare(K obj1, K obj2) {
        if (comparator == null) {
            if (obj1 instanceof Comparable && obj2 instanceof Comparable) {
                Comparable comp1 = (Comparable) obj1;
                Comparable comp2 = (Comparable) obj2;
                return comp1.compareTo(comp2);
            } else {
                throw new ClassCastException();
            }
        } else {
            return comparator.compare(obj1, obj2);
        }
    }

    private V removeEntry(SimpleEntry<K, V> entry) {
        /*
         doesn't actually remove the entry itself but mapping for the key
         stored i this entry
         */
        //assert entry != null;

        // find the next entry when traversing the tree symmetrically
        final V valueToReturn = entry.value;
        SimpleEntry<K, V> minEntryInRightSubtree = entry.right;
        if (minEntryInRightSubtree != null) {
            while (minEntryInRightSubtree.left != null) {
                minEntryInRightSubtree = minEntryInRightSubtree.left;
            }
            entry.key = minEntryInRightSubtree.key;
            entry.value = minEntryInRightSubtree.value;
            entry = minEntryInRightSubtree;
        }
        // found entry has no more than 1 son
        removeHalfLeafEntry(entry);
        --size;
        return valueToReturn;
    }

    private void removeHalfLeafEntry(SimpleEntry<K, V> entry) {
        SimpleEntry<K, V> sonToReplace = (entry.left != null
                ? entry.left : entry.right);

        if (sonToReplace != null) {
            sonToReplace.parent = entry.parent;
            if (entry.parent == null) {
                root = sonToReplace;
            } else if (entry == entry.parent.left) {
                entry.parent.left = sonToReplace;
            } else {
                entry.parent.right = sonToReplace;
            }
            entry.left = entry.right = entry.parent = null;

            if (entry.color == BLACK) {
                rebuildAfterRemoving(sonToReplace);
            }
        } else if (entry.parent == null) {
            root = null;
        } else { //  No children. 
            if (entry.color == BLACK) {
                rebuildAfterRemoving(entry);
            }

            if (entry.parent != null) {
                if (entry == entry.parent.left) {
                    entry.parent.left = null;
                } else if (entry == entry.parent.right) {
                    entry.parent.right = null;
                }
                entry.parent = null;
            }
        }
    }

    private void rebuildAfterRemoving(SimpleEntry entry) {
        if (entry != root && entry.color == BLACK) {
            if (entry == entry.parent.left) {
                SimpleEntry<K, V> sib = entry.parent.right;

                // case 3
                if (sib != null && sib.color == RED) {
                    sib.color = BLACK;
                    entry.parent.color = RED;
                    rotateLeft(entry.parent);
                    sib = entry.parent.right;
                }

                if (sib == null
                        || (sib.left == null || sib.left.color == BLACK)
                        && (sib.right == null || sib.right.color == BLACK)) {
                    // case 2
                    if (sib != null) {
                        sib.color = RED;
                    }
                    rebuildAfterRemoving(entry.parent);
                } else {
                    if (sib.right == null || sib.right.color == BLACK) {
                        sib.left.color = BLACK;
                        sib.color = RED;
                        rotateRight(sib);
                        sib = entry.parent.right;
                    }
                    sib.color = entry.parent.color;
                    entry.parent.color = BLACK;
                    sib.right.color = BLACK;
                    rotateLeft(entry.parent);
                }
            } else {
                SimpleEntry<K, V> sib = entry.parent.left;

                // case 3
                if (sib != null && sib.color == RED) {
                    sib.color = BLACK;
                    entry.parent.color = RED;
                    rotateRight(entry.parent);
                    sib = entry.parent.left;
                }

                if (sib == null
                        || (sib.right == null || sib.right.color == BLACK)
                        && (sib.left == null || sib.left.color == BLACK)) {
                    // case 2
                    if (sib != null) {
                        sib.color = RED;
                    }
                    rebuildAfterRemoving(entry.parent);
                } else {
                    if (sib.left == null || sib.left.color == BLACK) {
                        sib.right.color = BLACK;
                        sib.color = RED;
                        rotateLeft(sib);
                        sib = entry.parent.left;
                    }
                    sib.color = entry.parent.color;
                    entry.parent.color = BLACK;
                    sib.left.color = BLACK;
                    rotateRight(entry.parent);
                }
            }
        } else {
            entry.color = RED;
        }
    }

    public static class SimpleEntry<K, V> implements Entry<K, V> {

        static final boolean RED = true;
        static final boolean BLACK = false;

        K key;
        V value;
        SimpleEntry<K, V> parent;
        SimpleEntry<K, V> left;
        SimpleEntry<K, V> right;
        boolean color;

        public SimpleEntry(K key, V value, SimpleEntry<K, V> parent,
                SimpleEntry<K, V> left, SimpleEntry<K, V> right,
                boolean color) {
            this.key = key;
            this.value = value;
            this.parent = parent;
            this.left = left;
            this.right = right;
            this.color = color;
        }

        @Override
        public K getKey() {
            return key;
        }

        @Override
        public V getValue() {
            return value;
        }

        @Override
        public V setValue(V value) {
            final V oldValue = value;
            this.value = value;
            return oldValue;
        }
    }

    @Override
    public void clear() {
        unlinkTree(root);        //helps GC ?
        root = null;
        size = 0;
    }

    private void unlinkTree(SimpleEntry<K, V> tree) {
        if (tree != null) {
            tree.parent = null;
            unlinkTree(tree.left);
            unlinkTree(tree.right);
        }
    }

    /**
     * checks if this map contains a mapping with the specified key.
     *
     * @param key key to be checked
     * @return true if this map contains a mapping with the specified key, false
     * otherwise
     */
    @Override
    public boolean containsKey(K key) {
        return getEntry(key) != null;
    }

    /**
     * checks if this map contains a mapping with the specified value.
     *
     * @param value value to be checked
     * @return true if this map contains a mapping with the specified value,
     * false otherwise
     */
    @Override
    public boolean containsValue(V value) {
        if (root == null) {
            return false;
        } else {
            MyStack<SimpleEntry<K, V>> stack = new MyLinkedList<>();
            stack.push(root);
            while (true) {
                SimpleEntry<K, V> entry;
                try {
                    entry = stack.pop();
                } catch (MyNoSuchElementException ex) {
                    return false;
                }
                if (entry.getValue().equals(value)) {
                    return true;
                }
                if (entry.left != null) {
                    stack.push(entry.left);
                }
                if (entry.right != null) {
                    stack.push(entry.right);
                }
            }
        }
    }

    /**
     * Returns value associated with the specified key.
     *
     * @param key key value associated with which is to be returned
     * @return value associated with the specified key or {@code null} if there
     * is no such value
     */
    @Override
    public V get(K key) {
        final SimpleEntry<K, V> entry = getEntry(key);
        return (entry != null) ? entry.getValue() : null;
    }

    /**
     * Returns entry that corresponds the specified key if there is one, null
     * otherwise
     *
     * @param key
     * @return
     */
    private SimpleEntry<K, V> getEntry(K key) {
        return getEntry(key, root);
    }

    private SimpleEntry<K, V> getEntry(K key, SimpleEntry<K, V> tree) {
        if (tree == null) {
            compare(key, key);
            return null;
        } else {
            int compResult = compare(key, tree.getKey());
            if (compResult == 0) {
                return tree;
            } else if (compResult < 0) {
                return getEntry(key, tree.left);
            } else {
                return getEntry(key, tree.right);
            }
        }
    }

    /**
     * checks if this map is empty
     *
     * @return true if this map is empty, false otherwise
     */
    @Override
    public boolean isEmpty() {
        return size == 0;
    }

    /**
     * Returns the number of entries in this map.
     *
     * @return the number of entries in this map
     */
    @Override
    public int size() {
        return size;
    }

    /**
     * Returns the iterator over the entries in this map in symmetric order
     *
     * @return the iterator over the entries in this map in symmetric order
     */
    @Override
    public Iterator entryIterator() {
        return new EntryIterator();
    }

    private void rebuildAfterInsertion(SimpleEntry<K, V> entry) {
        SimpleEntry<K, V> father = entry.parent;
        if (father != null && father.color == RED) {
            // rule 4 is violated
            SimpleEntry<K, V> grandfather = father.parent;
            assert grandfather.color == BLACK;
            SimpleEntry<K, V> uncle;
            boolean leftGrandfather;
            boolean leftFather;
            if (grandfather.left == father) {
                uncle = grandfather.right;
                leftGrandfather = true;
            } else {
                uncle = grandfather.left;
                leftGrandfather = false;
            }
            leftFather = (entry == father.left);

            if (uncle != null && uncle.color == RED) { // case 2
                //    g - (B)                   g - (R)
                //     /   \                    /   \
                // f - (R) u - (R) -->    f - (B)  u - (B)
                //   /                     / 
                //i - (R)               i - (R)
                father.color = BLACK;
                uncle.color = BLACK;
                grandfather.color = RED;

                rebuildAfterInsertion(grandfather);
            } else { // case 1
                if (leftGrandfather) {
                    if (!leftFather) {
                        rotateLeft(father);
                        entry = father;
                        father = entry.parent;
                    }
                    rotateRight(grandfather);
                } else { // rightGrandFather
                    if (leftFather) {
                        rotateRight(father);
                        entry = father;
                        father = entry.parent;
                    }
                    rotateLeft(grandfather);
                }
                father.color = BLACK;
                grandfather.color = RED;
                entry.color = RED;
            }
        }
        root.color = BLACK;
    }

    private void rotateLeft(SimpleEntry<K, V> tree) {
        //       parent  (null)                           parent ( root )
        //            \                                     \    ( /    )
        //             tree                                  son
        //            /    \         ->                     /  \
        //           a      son                          tree    c
        //                 /  \                           / \
        //                b     c                        a   b
        SimpleEntry<K, V> son = tree.right;
        SimpleEntry<K, V> parent = tree.parent;

        tree.right = son.left;
        if (tree.right != null) {
            tree.right.parent = tree;
        }
        son.left = tree;
        son.parent = tree.parent;
        tree.parent = son;

        if (parent == null) {
            root = son;
        } else if (parent.left == tree) {
            parent.left = son;
        } else {
            parent.right = son;
        }
    }

    private void rotateRight(SimpleEntry<K, V> tree) {
        SimpleEntry<K, V> son = tree.left;
        SimpleEntry<K, V> parent = tree.parent;

        tree.left = son.right;
        if (tree.left != null) {
            tree.left.parent = tree;
        }
        son.right = tree;
        son.parent = tree.parent;
        tree.parent = son;

        if (parent == null) {
            root = son;
        } else if (parent.right == tree) {
            parent.right = son;
        } else {
            parent.left = son;
        }
    }

    private V insertToRight(K key, V value, SimpleEntry<K, V> tree) {
        tree.right = new SimpleEntry<>(key, value, tree, null, null, RED);
        ++size;
        rebuildAfterInsertion(tree.right);
        return null;
    }

    private V insertToLeft(K key, V value, SimpleEntry<K, V> tree) {
        tree.left = new SimpleEntry<>(key, value, tree, null, null, RED);
        ++size;
        rebuildAfterInsertion(tree.left);
        return null;
    }

    private V insert(K key, V value, SimpleEntry<K, V> tree) {
        if (tree == null) { // root
            compare(key, key); // check whether comparable and not null
            ++size;
            root = new SimpleEntry<>(key, value, null, null, null, BLACK);
            return null;
        } else {
            int compResult = compare(key, tree.getKey());
            if (compResult == 0) {
                final V v = tree.getValue();
                tree.setValue(value);
                return v;
            } else if (compResult < 0) {
                return (tree.left == null)
                        ? insertToLeft(key, value, tree)
                        : insert(key, value, tree.left);
            } else {
                return (tree.right == null)
                        ? insertToRight(key, value, tree)
                        : insert(key, value, tree.right);
            }
        }
    }

    class EntryIterator implements Iterator<SimpleEntry<K, V>> {

        MyStack<SimpleEntry<K, V>> entryStack = new MyLinkedList();

        SimpleEntry<K, V> next = root;

        SimpleEntry<K, V> lastReturned = null;

        public EntryIterator() {
        }

        @Override
        public boolean hasNext() {
            return (next != null || !((MyList) entryStack).isEmpty());
        }

        @Override
        public SimpleEntry<K, V> next() {
            while (next != null) {
                entryStack.push(next);
                next = next.left;
            }
            try {
                lastReturned = entryStack.pop();
                next = lastReturned.right;
                return lastReturned;
            } catch (MyNoSuchElementException ex) {
                return null;
            }
        }

        @Override
        public void remove() {
            if (lastReturned != null) {
                MyTreeMap.this.removeEntry(lastReturned);
                lastReturned = null;
            }
        }
    }

    // for debug only
//    public void log() {
//        log(root);
//        System.out.println();
//    }
//    
//    private void log(SimpleEntry<K, V> entry) {
//        if (entry != null) {
//            System.out.print(entry.key + ":" + entry.color);
//            System.out.print("(");
//            log(entry.left);
//            System.out.print("/");
//            log(entry.right);
//            System.out.print(")");
//        }
//    }
}
